import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time


def get_bin_reg_loss_approach(pred_reg, reg_label,distance_loc_scope, distance_bin_size,
                              new_width_scope,new_width_bin_size,
                              offset_center_scope,offset_center_bin_size,
                              azimuth_scope=360, azimuth_bin_size=60,
                              elevation_scope=75, elevation_bin_size=15,
                              width_scope=0.05, width_bin_size=0.01,
                              grasp_angle_scope=180, grasp_angle_bin_size=30):
    """
    Bin-based 3D bounding boxes regression loss. See https://arxiv.org/abs/1812.04244 for more details.

    :param pred_reg: (N, C)
    :param reg_label: (N, 3) [dx, dy, dz]
    :param loc_scope: constant
    :param loc_bin_size: constant
    :param num_head_bin: constant
    :param anchor_size: (N, 3) or (3)
    :param get_xz_fine:
    :param get_y_by_bin:
    :param loc_scope:
    :param loc_bin_size:
    :param get_ry_fine:
    :return:
    """

    reg_loss_dict = {}
    loc_loss = 0

    distance_offset_label = reg_label[:, 0]
    azimuth_offset_label, elevation_offset_label = reg_label[:, 1], reg_label[:, 2]
    width_offset_label = reg_label[:, 5]
    grasp_angle_offset_label = reg_label[:, 4]
    # new_width_offset_label = reg_label[:, 5]
    center_offset_distance_label =reg_label[:,6]

    per_loc_bin_num = int((distance_loc_scope[1] - distance_loc_scope[0]) / distance_bin_size)

    distance_shift = torch.clamp(distance_offset_label - distance_loc_scope[0], 0, distance_loc_scope[1] - distance_loc_scope[0] - 1e-4)
    distance_bin_label = (distance_shift / distance_bin_size).floor().long()

    distance_bin_l, distance_bin_r = 0, per_loc_bin_num

    loss_distance_bin = F.cross_entropy(pred_reg[:, distance_bin_l: distance_bin_r], distance_bin_label)
    reg_loss_dict['loss_distance_bin'] = loss_distance_bin.item()
    loc_loss += loss_distance_bin

    distance_res_l, distance_res_r = per_loc_bin_num , per_loc_bin_num * 2
    start_offset = distance_res_r
    distance_res_label = distance_shift - (distance_bin_label.float() * distance_bin_size + distance_bin_size / 2)
    distance_res_norm_label = distance_res_label / distance_bin_size
    distance_bin_onehot = torch.cuda.FloatTensor(distance_bin_label.size(0), per_loc_bin_num).zero_()
    distance_bin_onehot.scatter_(1, distance_bin_label.view(-1, 1).long(), 1)

    loss_distance_res = F.smooth_l1_loss((pred_reg[:, distance_res_l: distance_res_r] * distance_bin_onehot).sum(dim=1), distance_res_norm_label)
    reg_loss_dict['loss_distance_res'] = loss_distance_res.item()
    loc_loss += loss_distance_res

    per_azimuth_bin_num = int(azimuth_scope / azimuth_bin_size)
    per_elevation_bin_num = int(elevation_scope / elevation_bin_size)

    azimuth_bin_l, azimuth_bin_r = start_offset,  start_offset + per_azimuth_bin_num
    azimuth_res_l, azimuth_res_r = azimuth_bin_r, azimuth_bin_r + per_azimuth_bin_num
    start_offset = azimuth_res_r

    azimuth_shift = torch.clamp(azimuth_offset_label, 0, azimuth_scope - 1e-4)
    azimuth_bin_label = (azimuth_shift / azimuth_bin_size).floor().long()
    azimuth_res_label = azimuth_shift - (azimuth_bin_label.float() * azimuth_bin_size + azimuth_bin_size / 2)
    azimuth_res_norm_label = azimuth_res_label / azimuth_bin_size

    azimuth_bin_onehot = torch.cuda.FloatTensor(azimuth_bin_label.size(0), per_azimuth_bin_num).zero_()
    azimuth_bin_onehot.scatter_(1, azimuth_bin_label.view(-1, 1).long(), 1)

    loss_azimuth_bin = F.cross_entropy(pred_reg[:, azimuth_bin_l: azimuth_bin_r], azimuth_bin_label)
    loss_azimuth_res = F.smooth_l1_loss((pred_reg[:, azimuth_res_l: azimuth_res_r] * azimuth_bin_onehot).sum(dim=1),
                                        azimuth_res_norm_label)
    reg_loss_dict['loss_azimuth_bin'] = loss_azimuth_bin.item()
    reg_loss_dict['loss_azimuth_res'] = loss_azimuth_res.item()
    loc_loss += loss_azimuth_bin + loss_azimuth_res

    elevation_bin_l, elevation_bin_r = start_offset, start_offset + per_elevation_bin_num
    elevation_res_l, elevation_res_r = elevation_bin_r, elevation_bin_r + per_elevation_bin_num
    start_offset = elevation_res_r

    elevation_shift = torch.clamp(elevation_offset_label, 0, elevation_scope - 1e-4)
    elevation_bin_label = (elevation_shift / elevation_bin_size).floor().long()
    elevation_res_label = elevation_shift - (elevation_bin_label.float() * elevation_bin_size + elevation_bin_size / 2)
    elevation_res_norm_label = elevation_res_label / elevation_bin_size

    elevation_bin_onehot = torch.cuda.FloatTensor(elevation_bin_label.size(0), per_elevation_bin_num).zero_()
    elevation_bin_onehot.scatter_(1, elevation_bin_label.view(-1, 1).long(), 1)

    loss_elevation_bin = F.cross_entropy(pred_reg[:, elevation_bin_l: elevation_bin_r], elevation_bin_label)
    loss_elevation_res = F.smooth_l1_loss(
        (pred_reg[:, elevation_res_l: elevation_res_r] * elevation_bin_onehot).sum(dim=1), elevation_res_norm_label)
    reg_loss_dict['loss_elevation_bin'] = loss_elevation_bin.item()
    reg_loss_dict['loss_elevation_res'] = loss_elevation_res.item()

    loc_loss += loss_elevation_bin + loss_elevation_res

    per_width_bin_num = int(width_scope / width_bin_size)
    width_bin_l, width_bin_r = start_offset, start_offset + per_width_bin_num
    width_res_l, width_res_r = width_bin_r, width_bin_r + per_width_bin_num
    start_offset = width_res_r

    width_shift = torch.clamp(width_offset_label, 0, width_scope - 1e-4)
    width_bin_label = (width_shift / width_bin_size).floor().long()
    width_res_label = width_shift - (width_bin_label.float() * width_bin_size + width_bin_size / 2)
    width_res_norm_label = width_res_label / width_bin_size

    width_bin_onehot = torch.cuda.FloatTensor(width_bin_label.size(0), per_width_bin_num).zero_()
    width_bin_onehot.scatter_(1, width_bin_label.view(-1, 1).long(), 1)

    loss_width_bin = F.cross_entropy(pred_reg[:, width_bin_l: width_bin_r], width_bin_label)
    # width_res_norm_pred = torch.sigmoid(pred_reg[:, width_res_l: width_res_r] * width_bin_onehot).sum(dim=1) - 0.5
    width_res_norm_pred = (pred_reg[:, width_res_l: width_res_r] * width_bin_onehot).sum(dim=1)
    loss_width_res = F.smooth_l1_loss(width_res_norm_pred, width_res_norm_label)
    reg_loss_dict['loss_width_bin'] = loss_width_bin.item()
    reg_loss_dict['loss_width_res'] = loss_width_res.item()
    loc_loss += loss_width_bin + loss_width_res

    per_grasp_angle_bin_num = int(grasp_angle_scope / grasp_angle_bin_size)
    grasp_angle_bin_l, grasp_angle_bin_r = start_offset, start_offset + per_grasp_angle_bin_num
    grasp_angle_res_l, grasp_angle_res_r = grasp_angle_bin_r, grasp_angle_bin_r + per_grasp_angle_bin_num
    start_offset = grasp_angle_res_r

    grasp_angle_shift = torch.clamp(grasp_angle_offset_label, 0, grasp_angle_scope - 1e-4)
    grasp_angle_bin_label = (grasp_angle_shift / grasp_angle_bin_size).floor().long()
    grasp_angle_res_label = grasp_angle_shift - (
                grasp_angle_bin_label.float() * grasp_angle_bin_size + grasp_angle_bin_size / 2)
    grasp_angle_res_norm_label = grasp_angle_res_label / grasp_angle_bin_size

    grasp_angle_bin_onehot = torch.cuda.FloatTensor(grasp_angle_bin_label.size(0), per_grasp_angle_bin_num).zero_()
    grasp_angle_bin_onehot.scatter_(1, grasp_angle_bin_label.view(-1, 1).long(), 1)

    loss_grasp_angle_bin = F.cross_entropy(pred_reg[:, grasp_angle_bin_l: grasp_angle_bin_r], grasp_angle_bin_label)
    loss_grasp_angle_res = F.smooth_l1_loss(
        (pred_reg[:, grasp_angle_res_l: grasp_angle_res_r] * grasp_angle_bin_onehot).sum(dim=1),
        grasp_angle_res_norm_label)
    reg_loss_dict['loss_grasp_angle_bin'] = loss_grasp_angle_bin.item()
    reg_loss_dict['loss_grasp_angle_res'] = loss_grasp_angle_res.item()
    loc_loss += loss_grasp_angle_bin + loss_grasp_angle_res

    # per_new_width_bin_num = int(new_width_scope / new_width_bin_size)
    # new_width_bin_l, new_width_bin_r = start_offset, start_offset + per_new_width_bin_num
    # new_width_res_l, new_width_res_r = new_width_bin_r, new_width_bin_r + per_new_width_bin_num
    # start_offset = new_width_res_r

    # new_width_shift = torch.clamp(new_width_offset_label, 0, new_width_scope - 1e-4)
    # new_width_bin_label = (new_width_shift / new_width_bin_size).floor().long()
    # new_width_res_label = new_width_shift - (new_width_bin_label.float() * new_width_bin_size + new_width_bin_size / 2)
    # new_width_res_norm_label = new_width_res_label / new_width_bin_size
    #
    # new_width_bin_onehot = torch.cuda.FloatTensor(new_width_bin_label.size(0), per_new_width_bin_num).zero_()
    # new_width_bin_onehot.scatter_(1, new_width_bin_label.view(-1, 1).long(), 1)

    # loss_new_width_bin = F.cross_entropy(pred_reg[:, new_width_bin_l: new_width_bin_r], new_width_bin_label)
    # new_width_res_norm_pred = torch.sigmoid(pred_reg[:, new_width_res_l: new_width_res_r] * new_width_bin_onehot).sum(dim=1) - 0.5
    # loss_new_width_res = F.smooth_l1_loss(new_width_res_norm_pred, new_width_res_norm_label)
    # reg_loss_dict['loss_new_width_bin'] = loss_new_width_bin.item()
    # reg_loss_dict['loss_new_width_res'] = loss_new_width_res.item()
    # loc_loss += loss_new_width_bin + loss_new_width_res

    per_offset_center_bin_num = int(offset_center_scope / offset_center_bin_size) * 2
    offset_center_bin_l, offset_center_bin_r = start_offset, start_offset + per_offset_center_bin_num
    # offset_center_res_l,  offset_center_res_r = offset_center_bin_r, offset_center_bin_r + per_offset_center_bin_num

    offset_center_shift = torch.clamp(center_offset_distance_label + offset_center_scope, 0, offset_center_scope * 2 - 1e-4)
    offset_center_bin_label = (offset_center_shift / offset_center_bin_size).floor().long()
    # offset_center_res_label = offset_center_shift - (offset_center_bin_label.float() * offset_center_bin_size + offset_center_bin_size / 2)
    # offset_center_res_norm_label = offset_center_res_label / offset_center_bin_size

    offset_center_bin_onehot = torch.cuda.FloatTensor(offset_center_bin_label.size(0), per_offset_center_bin_num).zero_()
    offset_center_bin_onehot.scatter_(1, offset_center_bin_label.view(-1, 1).long(), 1)

    loss_offset_center_bin = F.cross_entropy(pred_reg[:, offset_center_bin_l: offset_center_bin_r], offset_center_bin_label)
    # offset_center_res_norm_pred = (pred_reg[:, offset_center_res_l: offset_center_res_r] * offset_center_bin_onehot).sum(dim=1)
    # loss_offset_center_res = F.smooth_l1_loss(offset_center_res_norm_pred, offset_center_res_norm_label)
    reg_loss_dict['loss_offset_center_bin'] = loss_offset_center_bin.item()
    # reg_loss_dict['loss_offset_center_res'] = loss_offset_center_res.item()
    loc_loss += loss_offset_center_bin #+ loss_offset_center_res

    reg_loss_dict['loss_loc'] = loc_loss

    return loc_loss, reg_loss_dict

